import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer

# 读取数据
data = pd.read_csv('../data/heart_disease.csv')

# 处理表格里面缺失的数据
data.dropna(inplace=True)
print('数据信息为', data.info())
print('数据头部信息', data.head())

# 2. 数据集的划分
X = data.drop('是否患有心脏病', axis=1)
y = data['是否患有心脏病']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)

# 3. 特征工程
# 数值型特征
numerical_features = ['年龄', '静息血压', '胆固醇', '最大心率', '运动后的ST下降', '主血管数量']
# 类别型特征
categorical_features = ['胸痛类型', '静息心电图结果', '峰值ST段的斜率', '地中海贫血']
# 二元特征
binary_features = ['性别', '空腹血糖', '运动性心绞痛']

# 创建列转换器
column_transformer = ColumnTransformer(transformers=[
    ('num', StandardScaler(), numerical_features),
    ('cat', OneHotEncoder(drop='first'), categorical_features),
    ('bin', 'passthrough', binary_features),
])

# 特征转换
X_train = column_transformer.fit_transform(X_train)
X_test = column_transformer.transform(X_test)

print(X_train.shape, X_test.shape)

# 4. 模型定义和训练
model = LogisticRegression()
model.fit(X_train, y_train)

print('模型准确率：', model.score(X_test, y_test))
